package com.beney.ds.trees.bst;

import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Queue;

/**
 * 方便起见，每一个结点只存储了key值，没有考虑value
 * 树的大部分问题都可以使用递归解决
 *
 * @param <K> 关键码的类型
 * @author Beney
 */
public class BST<K extends Comparable<K>> {
    Node root;

    public void setRoot(Node root) {
        this.root = root;
    }

    class Node {
        Node left, right;
        K key;

        public Node(K key) {
            this.key = key;
        }
    }

    //---------------------- 插入结点 ----------------------//
    public void insert(K key) {
        root = insert(root, key);
    }

    private Node insert(Node n, K key) {
        if (n == null) return new Node(key);
        int tmp;
        if ((tmp = key.compareTo(n.key)) < 0) {
            n.left = insert(n.left, key);
        } else if (tmp > 0) {
            n.right = insert(n.right, key);
        }
        return n;
    }

    //---------------- 删除结点 ------------------------//
    public void delete(K key) {
        root = delete(root, key);
    }

    private Node delete(Node n, K key) {
        if (n == null) return n;
        int tmp;
        if ((tmp = key.compareTo(n.key)) < 0) {
            n.left = delete(n.left, key);
        } else if (tmp > 0) {
            n.right = delete(n.right, key);
        } else {
            if (n.left == null)    // 只有一个孩子的情况 + 叶子的情况
                return n.right;
            else if (n.right == null)
                return n.left;
            Node oldN = n; // 有两个孩子的情况
            n = minNode(n.right); // 右子树中的最小结点替换为根
            n.right = deleteMin(oldN.right);
            n.left = oldN.left;
        }
        return n;
    }

    /**
     * @param x 子树根
     * @return 删除最小结点后的子树
     */
    private Node deleteMin(Node x) {
        if (x.left == null)
            return x.right;
        x.left = deleteMin(x.left);
        return x;
    }

    // n为根的子树中的最小结点
    private Node minNode(Node n) {
        Node tmp = null;
        while (n != null) {
            tmp = n;
            n = n.left;
        }
        return tmp;
    }


    //----------------- 层次遍历 ---------------------//
    public void levelTraversal() {
        if (root == null) return;
        Node tmp;
        Queue<Node> q = new LinkedList<>();
        q.add(root);
        while (!q.isEmpty()) {
            System.out.println((tmp = q.remove()).key);
            if (tmp.left != null) q.add(tmp.left);
            if (tmp.right != null) q.add(tmp.right);
        }
    }

    //------------------ 前序遍历 ----------------------//
    public void preOrder() {
        preOrder(root);
    }

    public void preOrder(Node r) {
        if (r == null) return;
        System.out.println(r.key);
        preOrder(r.left);
        preOrder(r.right);
    }

    //------------------ 中序遍历 ---------------------//
    public void inOrder() {
        inOrder(root);
    }

    public void inOrder(Node r) {
        if (r == null) return;
        inOrder(r.left);
        System.out.println(r.key);
        inOrder(r.right);
    }

    //----------------- 后序遍历 ---------------------//
    public void postOrder() {
        postOrder(root);
    }

    public void postOrder(Node r) {
        if (r == null) return;
        postOrder(r.left);
        postOrder(r.right);
        System.out.println(r.key);
    }


    /**
     * 根据前序遍历和中序遍历构造bst
     *
     * @param pre 前序遍历key值
     * @param in  中序遍历key值
     * @return 构建完的bst的根
     */
    public Node createBST1(List<K> pre, List<K> in) {
        if (pre == null || pre.size() == 0) return null;
        K k = pre.get(0);   // 前序遍历的第一个点key就是根节点key
        Node rt = new Node(k);
        int leftSize = in.indexOf(k);    // 根节点在中序遍历中的位置，也是左子树的大小
        rt.left = createBST1(
                pre.subList(1, 1 + leftSize),
                in.subList(0, leftSize));
        rt.right = createBST1(
                pre.subList(1 + leftSize, pre.size()),
                in.subList(leftSize + 1, pre.size()));
        return rt;
    }

    /**
     * 根据前序遍历，后序遍历构建bst
     *
     * @param pre  前序遍历
     * @param post 后序遍历
     * @return 构建完成的bst的根
     */
    public Node createBST2(List<K> pre, List<K> post) {
        if (pre == null || pre.size() == 0) return null;
        K k = pre.get(0);
        Node rt = new Node(k);
        Iterator<K> iter = pre.iterator();
        iter.next();        // 跳过当前根节点
        int leftSize = 0;   // 左子树大小
        while (iter.hasNext() && iter.next().compareTo(k) < 0)
            leftSize++;
        rt.left = createBST2(
                pre.subList(1, 1 + leftSize),
                post.subList(0, leftSize));
        rt.right = createBST2(
                pre.subList(leftSize + 1, pre.size()),
                post.subList(leftSize, pre.size() - 1));
        return rt;
    }

    /**
     * 求二叉树的节点数
     *
     * @return 二叉树的结点数
     */
    public int size() {
        return size(root);
    }

    /**
     * @param r 当前子树的根结点
     * @return r为根的子树的结点数
     */
    public int size(Node r) {
        if (r == null) return 0;
        return size(r.left) + size(r.right) + 1;
    }

    /**
     * @param k 层次(从1开始)
     * @return 第k层的结点个数
     */
    public int levelSize(int k) {
        return levelSize(k, 1, root);
    }

    private int levelSize(int k, int currentLevel, Node n) {
        if (n == null) return 0;
        else if (currentLevel == k) return 1;
        else if (currentLevel > k) return 0;
        return levelSize(k, currentLevel + 1, n.left) + levelSize(k, currentLevel + 1, n.right);
    }

    /**
     * @return bst的叶子数
     */
    public int leaves() {
        return leaves(root);
    }

    private int leaves(Node n) {
        if (n == null) return 0;
        else if (n.right == null && n.left == null) return 1;
        else
            return leaves(n.right) + leaves(n.left);
    }

    /**
     * @return bst的高度(最远的叶子到跟的路径长度)
     */
    public int height() {
        if (root == null) return 0;
        return height(root);
    }

    /**
     * @param n 当前子树的根结点
     * @return 结点n的高度
     */
    private int height(Node n) {
        if (n == null) return -1;
        return Math.max(height(n.left), height(n.right)) + 1;
    }


    /**
     * 判断当前二叉树是否是完全二叉树
     *
     * @return true：是完全二叉树
     */
    public boolean isComplete() {
        if (root == null) return true;
        Queue<Node> q = new LinkedList<>();
        q.add(root);
        Node n;
        while (!q.isEmpty()) {
            n = q.remove();
            if (n == null) {
                break;
            } else {
                q.add(n.left);
                q.add(n.right);
            }
        }
        while (!q.isEmpty()) {
            if (q.remove() != null)
                return false;
        }
        return true;
    }

    /**
     * 将二叉树镜像
     */
    public void mirror() {
        root = mirror(root);
    }

    private Node mirror(Node n) {
        if (n == null) return null;
        n.left = mirror(n.left);
        n.right = mirror(n.right);
        Node tmp = n.left;
        n.left = n.right;
        n.right = tmp;
        return n;
    }

    /**
     * @return 二叉树的直径
     */
    public int diameter() {
        return diameter(root);
    }

    private int diameter(Node n) {
        if (n == null) return 0;
        int leftHeight = height(n.left);
        int rightHeight = height(n.right);
        int leftDiameter = diameter(n.left);
        int rightDiameter = diameter(n.right);
        return Math.max(
                leftHeight + rightHeight + 2,
                Math.max(leftDiameter, rightDiameter)
        );
    }
}
